import torch
import onnx
import onnxruntime
import numpy as np

# 实现之前提到的 Debug 算子
# input: 张量 x 和调试名 name

# 这个算子“伪装”成一个普通的算子，使之能正常地参与推理、构建计算图的操作，
# 作为 torch.autograd.Function 的标准接口，ctx 必须作为参数保留
# 我们还是需要正确定义对输入 x 进行操作的 forward 函数
# 而在表示 PyTorch 与 ONNX 映射规则的 symbolic 函数里，
# 我们要定义一个带有调试名的 ONNX 算子，并把输入的 name 传给算子
class DebugOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, name):
        return x
    # 这里需要补充介绍算子定义函数 g.op() 的一些规范。
    # 在g.op()中，算子的属性需要以 {attibute_name}_{type}=attibute_value 这样的格式传入
    @staticmethod
    def symbolic(g, x, name):
        return g.op("my::Debug", x, name_s=name)
# 获取算子的调用接口。这样以后就可以通过 debug_apply(x, name) 来使用这个算子了
debug_apply = DebugOp.apply

# 实现精度对齐所需的所有操作
class Debugger():
    def __init__(self):
        super().__init__()
        self.torch_value = dict() # 记录了运行 PyTorch 模型后每个调试张量的值
        self.onnx_value = dict() # 记录了运行 ONNX 模型后每个调试张量的值
        self.output_debug_name = [] # 记录了把调试张量加入 ONNX 的输出后，每个输出张量的调试名

    def debug(self, x, name): # 封装了之前编写好的 debug_apply
        self.torch_value[name] = x.detach().cpu().numpy()
        # 在原 PyTorch 模型中调用，可以为导出的 ONNX 模型添加 Debug 算子节点
        return debug_apply(x, name)

    # 和 ONNX 的子模型提取函数的用法类似，
    # 可以把带调试节点的 ONNX 模型转化成一个可以输出调试张量的 ONNX 模型
    def extract_debug_model(self, input_path, output_path):
        model = onnx.load(input_path)
        inputs = [input.name for input in model.graph.input]
        outputs = []

        for node in model.graph.node:
            if node.op_type == 'Debug':
                debug_name = node.attribute[0].s.decode('ASCII')
                self.output_debug_name.append(debug_name)

                output_name = node.output[0]
                outputs.append(output_name)

                node.op_type = 'Identity'
                # 将当前节点的操作类型从 "Debug" 修改为 "Identity"，表示不执行任何操作，只传递输入
                # 为了消除 ONNX 不支持的 Debug 节点，一种比较简单的方式是直接把 Debug 节点修改成不执行任何操作的 Indentity 类型的节点。
                node.domain = ''
                del node.attribute[:]
        e = onnx.utils.Extractor(model)
        extracted = e.extract_model(inputs, outputs)
        onnx.save(extracted, output_path)

    def run_debug_model(self, input, debug_model):
        sess = onnxruntime.InferenceSession(debug_model,
                                            providers=['CPUExecutionProvider'])

        onnx_outputs = sess.run(None, input)
        for name, value in zip(self.output_debug_name, onnx_outputs):
            self.onnx_value[name] = value

    def print_debug_result(self):
        for name in self.torch_value.keys():
            if name in self.onnx_value:
                mse = np.mean((self.torch_value[name] - self.onnx_value[name])**2)
                print(f"{name} MSE: {mse}")

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.convs1 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1),
                                          torch.nn.Conv2d(3, 3, 3, 1, 1),
                                          torch.nn.Conv2d(3, 3, 3, 1, 1))
        self.convs2 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1),
                                          torch.nn.Conv2d(3, 3, 3, 1, 1))
        self.convs3 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1),
                                          torch.nn.Conv2d(3, 3, 3, 1, 1))
        self.convs4 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1),
                                          torch.nn.Conv2d(3, 3, 3, 1, 1),
                                          torch.nn.Conv2d(3, 3, 3, 1, 1))

    def forward(self, x):
        x = self.convs1(x)
        x = self.convs2(x)
        x = self.convs3(x)
        x = self.convs4(x)
        return x

torch_model = Model()
debugger = Debugger()

from types import MethodType


def new_forward(self, x):
    x = self.convs1(x)
    x = debugger.debug(x, 'x_0')
    x = self.convs2(x)
    x = debugger.debug(x, 'x_1')
    x = self.convs3(x)
    x = debugger.debug(x, 'x_2')
    x = self.convs4(x)
    x = debugger.debug(x, 'x_3')
    return x


torch_model.forward = MethodType(new_forward, torch_model)
# 为原模型新写一个 forward 函数。在这个新的函数函数中，我们可以通过 debugger.debug 把每一个输出张量标记起来，并各取一个不重复的调试名
dummy_input = torch.randn(1, 3, 10, 10)

# torch.onnx.export 模型使用的是跟踪法，模型的 new_forward 函数会被执行一次(===对应算子pytorch的值会被记录=)， debugger.debug 操作可以把 PyTorch 模型的调试张量输出记录在 debugger.torch_value
torch.onnx.export(torch_model, dummy_input, 'before_debug.onnx', input_names=['input'])

# 这一步会遍历所有graph，对节点名字为debug的算子作一些记录
debugger.extract_debug_model('before_debug.onnx', 'after_debug.onnx')

debugger.run_debug_model({'input':dummy_input.numpy()}, 'after_debug.onnx')
debugger.print_debug_result()



